#include <Arduino.h>
#include <ESP8266WiFi.h>
#include <WebSocketsClient.h>
#include <PubSubClient.h>
#include <time.h>
#include <EEPROM.h>
#include <ArduinoJson.h>
#include <WiFiUdp.h>

// // FluidNC WebSocket设置
// const char* wsHost = "192.168.102.100"; // FluidNC的IP地址
// const uint16_t wsPort = 81;           // WebSocket端口，通常是81

// MQTT设置
// const int mqtt_port = 8883;
// const char* mqtt_broker = "k702003e.ala.cn-hangzhou.emqxsl.cn";
// const char* mqtt_username = "emqx";
// const char* mqtt_password = "public";
String mqtt_to_ws_topic;    // MQTT发送到WebSocket的主题
String ws_to_mqtt_topic;    // WebSocket响应发送到MQTT的主题
String ws_status_topic;     // WebSocket状态主题


// NTP Server settings
const char *ntp_server = "cn.pool.ntp.org";     // 使用中国的NTP服务器
const long gmt_offset_sec = 8 * 3600;            // 设置为东八区 (UTC+8)
const int daylight_offset_sec = 0;               // 中国不使用夏令时


// SSL证书
static const char ca_cert[] PROGMEM = R"EOF(
-----BEGIN CERTIFICATE-----
MIIDrzCCApegAwIBAgIQCDvgVpBCRrGhdWrJWZHHSjANBgkqhkiG9w0BAQUFADBh
MQswCQYDVQQGEwJVUzEVMBMGA1UEChMMRGlnaUNlcnQgSW5jMRkwFwYDVQQLExB3
d3cuZGlnaWNlcnQuY29tMSAwHgYDVQQDExdEaWdpQ2VydCBHbG9iYWwgUm9vdCBD
QTAeFw0wNjExMTAwMDAwMDBaFw0zMTExMTAwMDAwMDBaMGExCzAJBgNVBAYTAlVT
MRUwEwYDVQQKEwxEaWdpQ2VydCBJbmMxGTAXBgNVBAsTEHd3dy5kaWdpY2VydC5j
b20xIDAeBgNVBAMTF0RpZ2lDZXJ0IEdsb2JhbCBSb290IENBMIIBIjANBgkqhkiG
9w0BAQEFAAOCAQ8AMIIBCgKCAQEA4jvhEXLeqKTTo1eqUKKPC3eQyaKl7hLOllsB
CSDMAZOnTjC3U/dDxGkAV53ijSLdhwZAAIEJzs4bg7/fzTtxRuLWZscFs3YnFo97
nh6Vfe63SKMI2tavegw5BmV/Sl0fvBf4q77uKNd0f3p4mVmFaG5cIzJLv07A6Fpt
43C/dxC//AH2hdmoRBBYMql1GNXRor5H4idq9Joz+EkIYIvUX7Q6hL+hqkpMfT7P
T19sdl6gSzeRntwi5m3OFBqOasv+zbMUZBfHWymeMr/y7vrTC0LUq7dBMtoM1O/4
gdW7jVg/tRvoSSiicNoxBN33shbyTApOB6jtSj1etX+jkMOvJwIDAQABo2MwYTAO
BgNVHQ8BAf8EBAMCAYYwDwYDVR0TAQH/BAUwAwEB/zAdBgNVHQ4EFgQUA95QNVbR
TLtm8KPiGxvDl7I90VUwHwYDVR0jBBgwFoAUA95QNVbRTLtm8KPiGxvDl7I90VUw
DQYJKoZIhvcNAQEFBQADggEBAMucN6pIExIK+t1EnE9SsPTfrgT1eXkIoyQY/Esr
hMAtudXH/vTBH1jLuG2cenTnmCmrEbXjcKChzUyImZOMkXDiqw8cvpOp/2PV5Adg
06O/nVsJ8dWO41P0jmP6P6fbtGbfYmbW0W5BjfIttep3Sp+dWOIrWcBAI+0tKIJF
PnlUkiaY4IBIqDfv8NZ5YBberOgOzW6sRBc4L0na4UU+Krk2U886UAb3LujEV0ls
YSEY1QSteDwsOoBrp+uvFRTp2InBuThs4pFsiv9kuXclVzDAGySj4dzp30d8tbQk
CAUw7C29C79Fv1C5qfPrmAESrciIxpg0X40KPMbp1ZWVbd4=
-----END CERTIFICATE-----
)EOF";

WebSocketsClient webSocket;
BearSSL::WiFiClientSecure espClient;
PubSubClient mqtt_client(espClient);

// 函数声明
void connectToMQTT();
bool syncTime();
void mqttCallback(char* topic, byte* payload, unsigned int length);

// 定义配置结构体
struct Config {
    // WiFi配置
    char wifi_ssid[32];
    char wifi_password[32];
    
    // FluidNC配置
    char ws_host[32];
    uint16_t ws_port;
    
    // MQTT配置
    char mqtt_broker[64];
    int mqtt_port;
    char mqtt_username[32];
    char mqtt_password[32];
};

Config config;

// EEPROM配置存储地址
#define CONFIG_ADDRESS 0
#define CONFIG_VERSION "v1"
#define CONFIG_SIZE 512

// 添加UDP相关变量
WiFiUDP udp;
const int UDP_PORT = 45678;  // 用于发现服务的UDP端口
const char* DISCOVERY_MESSAGE = "FLUIDNC_DISCOVER";
const unsigned long DISCOVERY_TIMEOUT = 5000;  // 5秒超时

// 添加自动发现函数
bool discoverFluidNC(char* host, uint16_t& port) {
    udp.begin(UDP_PORT);
    
    // 准备广播消息
    IPAddress broadcastIP(255, 255, 255, 255);
    udp.beginPacket(broadcastIP, UDP_PORT);
    udp.write(DISCOVERY_MESSAGE);
    udp.endPacket();
    
    Serial.println("正在搜索FluidNC设备...");
    
    unsigned long startTime = millis();
    while (millis() - startTime < DISCOVERY_TIMEOUT) {
        int packetSize = udp.parsePacket();
        if (packetSize) {
            char incomingPacket[255];
            int len = udp.read(incomingPacket, 255);
            if (len > 0) {
                incomingPacket[len] = 0;
                
                // 解析响应
                StaticJsonDocument<200> doc;
                DeserializationError error = deserializeJson(doc, incomingPacket);
                
                if (!error) {
                    if (doc.containsKey("type") && strcmp(doc["type"], "FLUIDNC") == 0) {
                        // 获取IP地址和端口
                        strlcpy(host, udp.remoteIP().toString().c_str(), 32);
                        port = doc["websocket_port"] | 81;
                        
                        Serial.printf("发现FluidNC设备: %s:%d\n", host, port);
                        udp.stop();
                        return true;
                    }
                }
            }
        }
        delay(100);
    }
    
    udp.stop();
    Serial.println("未发现FluidNC设备");
    return false;
}

// 保存配置到EEPROM
bool saveConfig() {
    EEPROM.begin(CONFIG_SIZE);
    
    // 清除EEPROM
    for (int i = 0; i < CONFIG_SIZE; i++) {
        EEPROM.write(i, 0);
    }
    
    // 使用JSON序列化配置
    StaticJsonDocument<CONFIG_SIZE> doc;
    
    doc["version"] = CONFIG_VERSION;
    doc["wifi_ssid"] = config.wifi_ssid;
    doc["wifi_password"] = config.wifi_password;
    doc["ws_host"] = config.ws_host;
    doc["ws_port"] = config.ws_port;
    doc["mqtt_broker"] = config.mqtt_broker;
    doc["mqtt_port"] = config.mqtt_port;
    doc["mqtt_username"] = config.mqtt_username;
    doc["mqtt_password"] = config.mqtt_password;
    
    // 序列化到字符串
    String jsonString;
    serializeJson(doc, jsonString);
    
    // 写入长度
    uint16_t length = jsonString.length();
    EEPROM.write(0, length & 0xFF);
    EEPROM.write(1, (length >> 8) & 0xFF);
    
    // 写入数据
    for (size_t i = 0; i < length; i++) {
        EEPROM.write(i + 2, jsonString[i]);
    }
    
    bool success = EEPROM.commit();
    if (success) {
        Serial.println("配置保存成功");
        Serial.println(jsonString);  // 打印保存的配置
    } else {
        Serial.println("配置保存失败！");
    }
    
    return success;  // 返回保存操作的结果
}

// 从EEPROM加载配置
bool loadConfig() {
    EEPROM.begin(CONFIG_SIZE);
    
    // 读取长度
    uint16_t length = EEPROM.read(0) | (EEPROM.read(1) << 8);
    
    if (length == 0 || length >= CONFIG_SIZE) {
        Serial.println("存储的配置无效");
        return false;
    }
    
    // 读取JSON字符串
    String jsonString;
    for (size_t i = 0; i < length; i++) {
        jsonString += (char)EEPROM.read(i + 2);
    }
    
    Serial.println("读取的配置：");
    Serial.println(jsonString);  // 打印读取的配置
    
    // 解析JSON
    StaticJsonDocument<CONFIG_SIZE> doc;
    DeserializationError error = deserializeJson(doc, jsonString);
    
    if (error) {
        Serial.print("JSON解析失败: ");
        Serial.println(error.c_str());
        return false;
    }
    
    // 检查版本
    const char* version = doc["version"] | "";
    if (strcmp(version, CONFIG_VERSION) != 0) {
        Serial.println("配置版本不匹配");
        return false;
    }
    
    // 加载配置
    strlcpy(config.wifi_ssid, doc["wifi_ssid"] | "", sizeof(config.wifi_ssid));
    strlcpy(config.wifi_password, doc["wifi_password"] | "", sizeof(config.wifi_password));
    strlcpy(config.ws_host, doc["ws_host"] | "", sizeof(config.ws_host));
    config.ws_port = doc["ws_port"] | 81;
    strlcpy(config.mqtt_broker, doc["mqtt_broker"] | "", sizeof(config.mqtt_broker));
    config.mqtt_port = doc["mqtt_port"] | 8883;
    strlcpy(config.mqtt_username, doc["mqtt_username"] | "", sizeof(config.mqtt_username));
    strlcpy(config.mqtt_password, doc["mqtt_password"] | "", sizeof(config.mqtt_password));
    
    Serial.println("当前配置：");
    Serial.printf("WiFi SSID: %s\n", config.wifi_ssid);
    Serial.printf("WebSocket Host: %s:%d\n", config.ws_host, config.ws_port);
    Serial.printf("MQTT Broker: %s:%d\n", config.mqtt_broker, config.mqtt_port);
    
    return true;
}

// 通过串口更新配置
void handleConfigCommand(const String& command) {
    Serial.print("收到配置命令: ");
    Serial.println(command);
    
    StaticJsonDocument<200> doc;
    DeserializationError error = deserializeJson(doc, command);
    
    if (error) {
        Serial.print("配置解析失败: ");
        Serial.println(error.c_str());
        return;
    }
    
    bool configChanged = false;
    
    // 更新配置
    if (doc.containsKey("wifi_ssid")) {
        strlcpy(config.wifi_ssid, doc["wifi_ssid"], sizeof(config.wifi_ssid));
        configChanged = true;
    }
    if (doc.containsKey("wifi_password")) {
        strlcpy(config.wifi_password, doc["wifi_password"], sizeof(config.wifi_password));
        configChanged = true;
    }
    
    // 更新WebSocket配置
    if (doc.containsKey("ws_host")) {
        strlcpy(config.ws_host, doc["ws_host"], sizeof(config.ws_host));
        configChanged = true;
    }
    if (doc.containsKey("ws_port")) {
        config.ws_port = doc["ws_port"].as<uint16_t>();
        configChanged = true;
    }
    
    // 更新MQTT配置
    if (doc.containsKey("mqtt_broker")) {
        strlcpy(config.mqtt_broker, doc["mqtt_broker"], sizeof(config.mqtt_broker));
        configChanged = true;
    }
    if (doc.containsKey("mqtt_port")) {
        config.mqtt_port = doc["mqtt_port"].as<int>();
        configChanged = true;
    }
    if (doc.containsKey("mqtt_username")) {
        strlcpy(config.mqtt_username, doc["mqtt_username"], sizeof(config.mqtt_username));
        configChanged = true;
    }
    if (doc.containsKey("mqtt_password")) {
        strlcpy(config.mqtt_password, doc["mqtt_password"], sizeof(config.mqtt_password));
        configChanged = true;
    }
    
    if (configChanged) {
        if (saveConfig()) {
            Serial.println("配置已保存，准备重启...");
            delay(1000);
            ESP.restart();
        }
    }
}

void webSocketEvent(WStype_t type, uint8_t * payload, size_t length) {
    String chip_id = String(ESP.getChipId(), HEX);
    String status_msg;

    switch(type) {
        case WStype_DISCONNECTED:
            status_msg = "WebSocket Disconnected! ChipID: " + chip_id;
            Serial.println("WebSocket断开连接");
            mqtt_client.publish(ws_status_topic.c_str(), status_msg.c_str());
            break;
            
        case WStype_CONNECTED:
            status_msg = "WebSocket Connected! ChipID: " + chip_id;
            Serial.println("WebSocket已连接");
            mqtt_client.publish(ws_status_topic.c_str(), status_msg.c_str());
            webSocket.sendTXT("?");
            break;
            
        case WStype_ERROR:
            status_msg = "WebSocket Error! ChipID: " + chip_id + 
                        ", Host: " + String(config.ws_host) + 
                        ", Port: " + String(config.ws_port);
            Serial.println("WebSocket发生错误");
            mqtt_client.publish(ws_status_topic.c_str(), status_msg.c_str());
            break;
            
        case WStype_PING:
            // Serial.println("收到Ping");
            // mqtt_client.publish(ws_status_topic.c_str(), status_msg.c_str());

            // mqtt_client.publish(ws_to_mqtt_topic, "WebSocket Ping received");
            break;
            
        case WStype_PONG:
            // Serial.println("收到Pong");

            // mqtt_client.publish(ws_to_mqtt_topic, "WebSocket Pong received");
            break;

        case WStype_BIN:
            if (length > 0) {
                char* message = new char[length + 1];
                memcpy(message, payload, length);
                message[length] = '\0';
                
                mqtt_client.publish(ws_to_mqtt_topic.c_str(), message);
                
                // 按行处理消息并输出到串口
                char* line = strtok(message, "\n");
                while (line != NULL) {
                    Serial.println(line);
                    line = strtok(NULL, "\n");
                }
                
                delete[] message;
            }
            break;
            
        case WStype_FRAGMENT_TEXT_START:
        case WStype_FRAGMENT_BIN_START:
            // status_msg = "WebSocket Fragment Start - ChipID: " + chip_id;
            // mqtt_client.publish(ws_to_mqtt_topic, status_msg.c_str());
            break;
            
        case WStype_FRAGMENT:
            // status_msg = "WebSocket Fragment Continue - ChipID: " + chip_id;
            // mqtt_client.publish(ws_to_mqtt_topic, status_msg.c_str());
            break;
            
        case WStype_FRAGMENT_FIN:
            // status_msg = "WebSocket Fragment End - ChipID: " + chip_id;
            // mqtt_client.publish(ws_to_mqtt_topic, status_msg.c_str());
            break;
            
        default:
            // status_msg = "WebSocket Unknown Event: " + String(type) + 
            //             " - ChipID: " + chip_id;
            // mqtt_client.publish(ws_status_topic.c_str(), status_msg.c_str());
            break;
    }
}

// 修改WebSocket连接状态检查函数
void checkWebSocketConnection() {
    static unsigned long lastCheck = 0;
    static int reconnectAttempts = 0;
    const int maxReconnectAttempts = 5;
    
    if (millis() - lastCheck > 5000) {
        lastCheck = millis();
        
        if (webSocket.sendPing()) {
            reconnectAttempts = 0;
        } else {
            String status_msg = "WebSocket Connection Failed! ChipID: " + 
                              String(ESP.getChipId(), HEX) + 
                              ", Target: " + String(config.ws_host) + ":" + 
                              String(config.ws_port);
            
            Serial.println(status_msg);
            mqtt_client.publish(ws_status_topic.c_str(), status_msg.c_str());
            
            reconnectAttempts++;
            if (reconnectAttempts >= maxReconnectAttempts) {
                // 尝试重新发现设备
                if (discoverFluidNC(config.ws_host, config.ws_port)) {
                    saveConfig();
                    webSocket.begin(config.ws_host, config.ws_port, "/");
                    reconnectAttempts = 0;
                } else {
                    status_msg = "WebSocket Max Reconnect Attempts Reached! Restarting...";
                    mqtt_client.publish(ws_status_topic.c_str(), status_msg.c_str());
                    delay(1000);
                    ESP.restart();
                }
            }
        }
    }
}

void mqttCallback(char *topic, byte *payload, unsigned int length) {
    if (String(topic) == mqtt_to_ws_topic) {  // 使用String比较
        // 创建临时缓冲区
        char* command = new char[length + 1];
        memcpy(command, payload, length);
        command[length] = '\0';
        
        // 处理命令
        if (length == 1) {
            // 单字符命令直接发送
            webSocket.sendTXT(command);
        } else {
            // 其他命令添加换行符
            String fullCommand = String(command) + "\n";
            webSocket.sendTXT(fullCommand);
        }
        
        delete[] command;
    }
}


bool syncTime() {
    configTime(gmt_offset_sec, daylight_offset_sec, ntp_server, "time.nist.gov", "time.windows.com");
    Serial.print("Waiting for NTP time sync: ");
    
    int retry = 0;
    const int maxRetries = 20;  // 增加等待时间
    time_t now = 0;
    
    while ((now = time(nullptr)) < 8 * 3600 * 2) {
        delay(500);
        Serial.print(".");
        retry++;
        if (retry >= maxRetries) {
            Serial.println("\nNTP sync timeout!");
            return false;
        }
    }
    
    struct tm* timeinfo = localtime(&now);
    if (timeinfo) {
        Serial.print("\nCurrent time: ");
        Serial.println(asctime(timeinfo));
        return true;
    } else {
        Serial.println("\nFailed to obtain local time");
        return false;
    }
}

void connectToMQTT() {
    BearSSL::X509List serverTrustedCA(ca_cert);
    espClient.setTrustAnchors(&serverTrustedCA);
    espClient.setX509Time(time(nullptr));
    
    int retries = 0;
    while (!mqtt_client.connected() && retries < 5) {
        // 获取芯片ID
        String chip_id = String(ESP.getChipId(), HEX);
        String client_id = "esp8266-" + chip_id;
        
        Serial.printf("芯片ID: %s\n", chip_id.c_str());
        Serial.printf("连接MQTT服务器 %s...\n", client_id.c_str());
        
        // 使用配置中的用户名和密码
        if (mqtt_client.connect(client_id.c_str(), config.mqtt_username, config.mqtt_password)) {
            Serial.println("MQTT已连接");
            mqtt_client.subscribe(mqtt_to_ws_topic.c_str());  // 订阅命令主题
            
            // 发布欢迎消息到状态主题
            String welcome_msg = "ESP8266 Online! ChipID: " + chip_id + 
                               ", IP: " + WiFi.localIP().toString();
            mqtt_client.publish(ws_status_topic.c_str(), welcome_msg.c_str());
            
            return;
        }
        
        retries++;
        delay(5000);
    }
}

void setup() {
    Serial.begin(115200);
    Serial.println("\n\n启动中...");
    randomSeed(micros());

    // 设置所有主题（包含芯片ID）
    String chip_id = String(ESP.getChipId(), HEX);
    mqtt_to_ws_topic = "fluidnc/" + chip_id + "/commands";
    ws_to_mqtt_topic = "fluidnc/" + chip_id + "/responses";
    ws_status_topic = "fluidnc/" + chip_id + "/status";
    
    // 加载配置
    if (!loadConfig()) {
        Serial.println("使用默认配置");
        // 设置默认配置
        strlcpy(config.wifi_ssid, "hzg", sizeof(config.wifi_ssid));
        strlcpy(config.wifi_password, "12345678", sizeof(config.wifi_password));
        strlcpy(config.ws_host, "192.168.0.119", sizeof(config.ws_host));
        config.ws_port = 81;
        strlcpy(config.mqtt_broker, "k702003e.ala.cn-hangzhou.emqxsl.cn", sizeof(config.mqtt_broker));
        config.mqtt_port = 8883;
        strlcpy(config.mqtt_username, "emqx", sizeof(config.mqtt_username));
        strlcpy(config.mqtt_password, "public", sizeof(config.mqtt_password));
        
        // 保存默认配置
        saveConfig();
    }

    // 设置WiFi模式
    WiFi.mode(WIFI_STA);
    
    // 使用配置中的WiFi信息进行连接
    Serial.printf("正在连接到WiFi: %s\n", config.wifi_ssid);
    WiFi.begin(config.wifi_ssid, config.wifi_password);
    
    int attempts = 0;
    while (WiFi.status() != WL_CONNECTED && attempts < 20) {
        delay(500);
        Serial.print(".");
        attempts++;
    }
    
    if (WiFi.status() == WL_CONNECTED) {
        Serial.printf("\nWiFi已连接，IP地址: %s\n", WiFi.localIP().toString().c_str());
        
        // 尝试自动发现FluidNC设备
        if (discoverFluidNC(config.ws_host, config.ws_port)) {
            // 发现设备后保存配置
            saveConfig();
        } else {
            Serial.println("使用已保存的WebSocket配置");
        }
        
        // 设置WebSocket连接
        webSocket.begin(config.ws_host, config.ws_port, "/");
        webSocket.onEvent(webSocketEvent);
        webSocket.setReconnectInterval(5000);
    } else {
        Serial.println("\nWiFi连接失败！");
        // 连接失败时不要立即重启，给用户机会通过串口修改配置
        Serial.println("请通过串口发送新的WiFi配置");
        while (true) {
            if (Serial.available()) {
                String command = Serial.readStringUntil('\n');
                handleConfigCommand(command);
                // 配置更新后重启
                ESP.restart();
            }
            delay(100);
        }
    }

    // 同步时间
    // if (!syncTime()) {
    //     ESP.restart();
    // }

    
    // 多次尝试同步时间
    bool timeSync = false;
    for(int i = 0; i < 3; i++) {
        if(syncTime()) {
            timeSync = true;
            break;
        }
        delay(1000);
    }

    

    // 设置MQTT
    mqtt_client.setServer(config.mqtt_broker, config.mqtt_port);
    mqtt_client.setCallback(mqttCallback);
    connectToMQTT();
}

void loop() {
    webSocket.loop();
    mqtt_client.loop();

    // 检查WiFi连接状态
    if (WiFi.status() != WL_CONNECTED) {
        String status_msg = "WiFi Disconnected! ChipID: " + String(ESP.getChipId(), HEX);
        mqtt_client.publish(ws_status_topic.c_str(), status_msg.c_str());
        Serial.println("WiFi连接断开，重启设备...");
        delay(1000);
        ESP.restart();
    }

    // 检查MQTT连接
    if (!mqtt_client.connected()) {
        connectToMQTT();
    }

    // 检查WebSocket连接状态
    checkWebSocketConnection();

    // 处理串口输入
    if (Serial.available()) {
        String command = Serial.readStringUntil('\n');
        command.trim();
        
        if (command.length() == 1) {
            webSocket.sendTXT(command);
        } else {
            command += "\n";
            webSocket.sendTXT(command);
        }
    }

    delay(10);
}
